"""
Pinecone 手写数字识别示例
用于演示如何使用 Pinecone 进行向量相似度搜索来识别手写数字
"""

import numpy as np
from sklearn.datasets import load_digits
from pinecone import Pinecone, ServerlessSpec
import os
import time

# 初始化 Pinecone
# 请在环境变量中设置 PINECONE_API_KEY
# 或者直接替换下面的 os.getenv("PINECONE_API_KEY") 为你的 API Key
API_KEY = os.getenv("PINECONE_API_KEY", "pcsk_C1E8x_PBoUu1QoydYsUK4K5SWe5EDQvNskevRHiCpc7n8ozSny7Axtf99Ewm2F9LR9F7S")

print("=" * 60)
print("Pinecone 手写数字识别示例")
print("=" * 60)
print("\n如果你还没有 Pinecone 账号，请按照以下步骤操作：")
print("1. 访问 https://www.pinecone.io/")
print("2. 注册免费账号")
print("3. 在控制台中创建 API Key")
print("4. 将 API Key 设置为环境变量 PINECONE_API_KEY")
print("   或直接在本文件中替换 API_KEY 的值")
print("\n" + "=" * 60 + "\n")

# 初始化 Pinecone 客户端
pc = Pinecone(api_key=API_KEY)

# 索引配置
INDEX_NAME = "mnist-digits-example"
DIMENSION = 64  # MNIST digits 数据集每个样本 8x8=64 维

# 检查索引是否存在，如果不存在则创建
existing_indexes = [index.name for index in pc.list_indexes()]

if INDEX_NAME not in existing_indexes:
    print(f"创建新索引: {INDEX_NAME}")
    pc.create_index(
        name=INDEX_NAME,
        dimension=DIMENSION,
        metric="euclidean",
        spec=ServerlessSpec(
            cloud="aws",
            region="us-east-1"
        )
    )
    # 等待索引创建完成
    while not pc.describe_index(INDEX_NAME).status['ready']:
        print("等待索引创建完成...")
        time.sleep(1)
    print("索引创建完成！")
else:
    print(f"索引 {INDEX_NAME} 已存在")

# 连接到索引
index = pc.Index(INDEX_NAME)

# 加载 MNIST digits 数据集
print("\n加载 MNIST digits 数据集...")
digits = load_digits()
X = digits.data  # shape: (1797, 64)
y = digits.target  # shape: (1797,)

# 准备一些示例数据上传到 Pinecone（如果索引为空）
stats = index.describe_index_stats()
if stats['total_vector_count'] == 0:
    print("\n上传示例数据到 Pinecone...")
    # 只上传前 100 个样本作为示例
    vectors_to_upsert = []
    for i in range(100):
        vectors_to_upsert.append({
            "id": f"digit_{i}",
            "values": X[i].tolist(),
            "metadata": {"label": int(y[i])}
        })
    
    # 批量上传
    index.upsert(vectors=vectors_to_upsert)
    print(f"成功上传 {len(vectors_to_upsert)} 条数据")
    # 等待索引更新
    print("等待索引更新...")
    time.sleep(5)
else:
    print(f"\n索引中已有 {stats['total_vector_count']} 条数据")

# 进行查询测试
print("\n" + "=" * 60)
print("进行查询测试")
print("=" * 60)

# 选择一个测试样本（第 50 个样本，在上传的前 100 个中）
test_idx = 50
test_vector = X[test_idx].tolist()
test_label = int(y[test_idx])

print(f"\n测试样本索引: {test_idx}")
print(f"真实标签: {test_label}")

# 查询最相似的 k=11 个向量
k = 11
query_results = index.query(
    vector=test_vector,
    top_k=k,
    include_metadata=True
)

print(f"\n查询结果 (k={k}):")
print("-" * 60)

# 收集所有邻居的标签用于投票
neighbor_labels = []
for i, match in enumerate(query_results['matches']):
    match_id = match['id']
    distance = match['score']  # 对于 euclidean，score 就是距离
    label = match['metadata']['label']
    neighbor_labels.append(label)
    print(f"{i+1}. id: {match_id}, distance: {distance:.4f}, label: {label}")

# 使用多数投票预测
if neighbor_labels:
    predicted_digit = max(set(neighbor_labels), key=neighbor_labels.count)
    
    print("-" * 60)
    print(f"\n预测结果: {predicted_digit}")
    print(f"真实标签: {test_label}")
    print(f"预测{'正确' if predicted_digit == test_label else '错误'}！")
else:
    print("警告：查询结果为空！")
    print("可能原因：索引还没有完全准备好，请稍后再试。")

print("\n" + "=" * 60)
print("示例完成！")
print("=" * 60)
print("\n提示：")
print("- 这个示例展示了如何使用 Pinecone 进行 KNN 搜索")
print("- 你可以运行 pinecone_train.py 来使用完整的训练集")
print("- 你可以运行 optimal_knn_webapp_pinecone.py 来启动 Web 应用")
